import pandas as pd
import numpy as np
import pickle as pk
import seaborn as sns
import matplotlib.pyplot as plt
import math
import scipy.io as sio
import scipy.sparse as ss

data_path = '../Data/'  # 文件路径
model_path = '../model/' # 模型路径

with open(model_path+'data_all_train.pkl', 'rb') as fr:
    data_all_train = pk.load(fr)
fr.close()

def artist_name_2_set(item_name, item_set):
    item_name = str(item_name)
    item_arr = item_name.split('|')
    for item in item_arr:
        item = item.strip()
        item_set.add(item)
    return

artist_name_set = set()
data_all_train.artist_name.apply(lambda artist_name : artist_name_2_set(artist_name, artist_name_set))
artist_name_list = list(artist_name_set)
artist_name_len = len(artist_name_list)
print('artist_name_len: %d ' % artist_name_len)
# 建立 genre_id 索引
artist_name_index = dict()
for i,artist_name in enumerate(artist_name_list):
    artist_name_index[artist_name] = i


def get_item_matrix(data_all_train, item_index_dict, item_len):
    item_matrix = ss.dok_matrix((data_all_train.shape[0], item_len), dtype='uint8')
    for df_index in range(data_all_train.shape[0]):
        line = data_all_train.iloc[df_index, :]
        for item in str(line.artist_name).split('|'):
             item = item.strip()
             item_index = item_index_dict[item]
             item_matrix[df_index,item_index] = 1
    return item_matrix

print('get_item_matrix ...')
artist_name_matrix = get_item_matrix(data_all_train,artist_name_index,artist_name_len)
print('get_item_matrix end')
# 生成 artist_name 在 dataframe 中的列名
artist_name_columns = list()
for i in range(artist_name_len):
    artist_name_columns.append('artist_name_'+str(i))

try:
    sio.mmwrite(model_path+'artist_name_matrix.mtx', artist_name_matrix)
except Exception as e:
    print('sio exception:',e)

try:    
    artist_name_df = pd.DataFrame(artist_name_matrix.todense(), columns = artist_name_columns, index = data_all_train.index,dtype='uint8')  
except Exception as e:
    print('artist_name_df exception:',e)

try:
    print('artist_name_df shape: %d, %d' % artist_name_df.shape)
except Exception as e:
    print(str(e))
    
# artist_name_df 计算该 dataframe 占用内存大小
artist_name_df_mem = artist_name_df.memory_usage().sum()/(1024**2)
print('artist_name_df_mem memory usage: %f Mb' % artist_name_df_mem)
try:
    # 序列化到磁盘中
    with open(model_path + 'artist_name_df.pkl','wb') as fw:
        pk.dump(artist_name_df,fw)
    fw.close()
except Exception as e:
    print('failed to save artist_name_df',e)


